"""
Smart Memory Manager Node for ComfyUI

Provides intelligent, automated memory management combining all features.
"""

import time
import threading
from ..utils.memory_utils import (
    get_memory_info, cleanup_memory, optimize_vram, check_memory_pressure,
    format_bytes, global_memory_manager, MemoryTracker
)


class SmartMemoryManagerNode:
    """
    Smart Memory Manager Node - Intelligent automated memory management
    """
    
    def __init__(self):
        self.tracker = MemoryTracker()
        self.management_active = False
        self.management_thread = None
        self.stats = {
            'cleanups_performed': 0,
            'vram_optimizations': 0,
            'memory_freed_total': 0,
            'leaks_detected': 0,
            'uptime_start': time.time()
        }
        self.last_action = "Initialized"
        
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "enable_smart_management": ("BOOLEAN", {"default": True}),
                "management_mode": (["Conservative", "Balanced", "Aggressive"], {
                    "default": "Balanced"
                }),
                "memory_warning_threshold": ("FLOAT", {
                    "default": 75.0,
                    "min": 60.0,
                    "max": 90.0,
                    "step": 5.0,
                    "display": "slider"
                }),
                "memory_critical_threshold": ("FLOAT", {
                    "default": 85.0,
                    "min": 70.0,
                    "max": 95.0,
                    "step": 5.0,
                    "display": "slider"
                }),
            },
            "optional": {
                "check_interval": ("FLOAT", {
                    "default": 45.0,
                    "min": 15.0,
                    "max": 300.0,
                    "step": 15.0,
                    "display": "slider"
                }),
                "enable_leak_detection": ("BOOLEAN", {"default": True}),
                "enable_vram_optimization": ("BOOLEAN", {"default": True}),
                "auto_cleanup_aggressive": ("BOOLEAN", {"default": False}),
                "leak_detection_interval": ("FLOAT", {
                    "default": 300.0,  # 5 minutes
                    "min": 60.0,
                    "max": 1800.0,  # 30 minutes
                    "step": 60.0,
                    "display": "slider"
                }),
            }
        }
    
    RETURN_TYPES = (
        "STRING",      # Management status
        "STRING",      # Detailed report
        "BOOLEAN",     # Management active
        "FLOAT",       # Current memory usage
        "STRING",      # Last action taken
        "STRING",      # Performance statistics
    )
    
    RETURN_NAMES = (
        "management_status",
        "detailed_report",
        "management_active",
        "current_memory_usage",
        "last_action",
        "performance_stats",
    )
    
    FUNCTION = "manage_memory_smart"
    CATEGORY = "Memory Management"
    OUTPUT_NODE = True
    
    def manage_memory_smart(self, enable_smart_management, management_mode, 
                          memory_warning_threshold, memory_critical_threshold,
                          check_interval=45.0, enable_leak_detection=True,
                          enable_vram_optimization=True, auto_cleanup_aggressive=False,
                          leak_detection_interval=300.0):
        """Intelligent memory management with comprehensive monitoring"""
        
        try:
            # Update configuration
            config = {
                'mode': management_mode,
                'warning_threshold': memory_warning_threshold,
                'critical_threshold': memory_critical_threshold,
                'check_interval': check_interval,
                'leak_detection': enable_leak_detection,
                'vram_optimization': enable_vram_optimization,
                'aggressive_cleanup': auto_cleanup_aggressive,
                'leak_interval': leak_detection_interval
            }
            
            # Start or stop management
            if enable_smart_management and not self.management_active:
                self._start_smart_management(config)
                action = "Smart management started"
            elif not enable_smart_management and self.management_active:
                self._stop_smart_management()
                action = "Smart management stopped"
            elif enable_smart_management and self.management_active:
                self._update_configuration(config)
                action = "Configuration updated"
            else:
                action = "Management inactive"
            
            self.last_action = action
            
            # Get current memory state
            memory_info = get_memory_info()
            current_usage = memory_info.get('system_percent', 0)
            
            # Create reports
            status_report = self._create_status_report(enable_smart_management, config, current_usage)
            detailed_report = self._create_detailed_report(memory_info, config)
            performance_stats = self._create_performance_stats()
            
            return (
                status_report,
                detailed_report,
                self.management_active,
                current_usage,
                action,
                performance_stats
            )
            
        except Exception as e:
            error_msg = f"Smart memory management error: {str(e)}"
            return (
                error_msg,
                error_msg,
                False,
                0.0,
                "Error occurred",
                "No statistics available"
            )
    
    def _start_smart_management(self, config):
        """Start smart memory management"""
        
        if self.management_active:
            return
        
        self.management_active = True
        self.stats['uptime_start'] = time.time()
        
        # Start leak detection if enabled
        if config['leak_detection']:
            self.tracker.start_tracking()
            self.tracker.take_snapshot("Smart_Manager_Baseline")
        
        # Start management thread
        self.management_thread = threading.Thread(
            target=self._smart_management_loop,
            args=(config,),
            daemon=True
        )
        self.management_thread.start()
    
    def _stop_smart_management(self):
        """Stop smart memory management"""
        
        self.management_active = False
        
        if self.management_thread:
            self.management_thread.join(timeout=10.0)
        
        # Stop leak detection
        if hasattr(self.tracker, 'stop_tracking'):
            self.tracker.stop_tracking()
    
    def _update_configuration(self, config):
        """Update management configuration"""
        
        # Configuration is automatically picked up by the management loop
        # through the config parameter passed to the thread
        pass
    
    def _smart_management_loop(self, config):
        """Main smart management loop"""
        
        last_leak_check = time.time()
        snapshot_counter = 0
        
        while self.management_active:
            try:
                # Get current memory state
                memory_info = get_memory_info()
                memory_percent = memory_info.get('system_percent', 0)
                pressure_info = check_memory_pressure()
                
                # Determine action based on memory usage and mode
                action_taken = False
                
                # Critical threshold - immediate action required
                if memory_percent >= config['critical_threshold']:
                    self._perform_critical_cleanup(config)
                    action_taken = True
                    self.last_action = f"Critical cleanup (memory at {memory_percent:.1f}%)"
                
                # Warning threshold - preventive action
                elif memory_percent >= config['warning_threshold']:
                    self._perform_warning_cleanup(config)
                    action_taken = True
                    self.last_action = f"Warning cleanup (memory at {memory_percent:.1f}%)"
                
                # Periodic maintenance based on mode
                elif config['mode'] == 'Aggressive':
                    if time.time() % (config['check_interval'] * 2) < config['check_interval']:
                        self._perform_maintenance(config)
                        action_taken = True
                        self.last_action = "Aggressive maintenance"
                
                # Leak detection
                current_time = time.time()
                if (config['leak_detection'] and 
                    current_time - last_leak_check >= config['leak_interval']):
                    
                    self._perform_leak_detection()
                    last_leak_check = current_time
                
                # Sleep until next check
                time.sleep(config['check_interval'])
                
            except Exception as e:
                print(f"Smart management error: {e}")
                time.sleep(config['check_interval'])
    
    def _perform_critical_cleanup(self, config):
        """Perform critical memory cleanup"""
        
        # Aggressive cleanup
        cleanup_result = cleanup_memory(aggressive=True)
        
        # VRAM optimization if enabled
        if config['vram_optimization']:
            vram_result = optimize_vram()
            self.stats['vram_optimizations'] += 1
        
        # Update statistics
        self.stats['cleanups_performed'] += 1
        self.stats['memory_freed_total'] += cleanup_result.get('memory_freed', 0)
    
    def _perform_warning_cleanup(self, config):
        """Perform warning-level cleanup"""
        
        # Standard cleanup
        aggressive = config['aggressive_cleanup'] or config['mode'] == 'Aggressive'
        cleanup_result = cleanup_memory(aggressive=aggressive)
        
        # VRAM optimization if enabled and mode is not conservative
        if config['vram_optimization'] and config['mode'] != 'Conservative':
            vram_result = optimize_vram()
            self.stats['vram_optimizations'] += 1
        
        # Update statistics
        self.stats['cleanups_performed'] += 1
        self.stats['memory_freed_total'] += cleanup_result.get('memory_freed', 0)
    
    def _perform_maintenance(self, config):
        """Perform routine maintenance"""
        
        # Light cleanup
        cleanup_result = cleanup_memory(aggressive=False)
        
        # Light VRAM optimization
        if config['vram_optimization']:
            vram_result = optimize_vram()
        
        # Update statistics
        self.stats['memory_freed_total'] += cleanup_result.get('memory_freed', 0)
    
    def _perform_leak_detection(self):
        """Perform leak detection"""
        
        try:
            # Take snapshot
            snapshot = self.tracker.take_snapshot(f"Smart_Check_{int(time.time())}")
            
            if len(self.tracker.snapshots) >= 2:
                leaks = self.tracker.detect_leaks()
                if leaks:
                    self.stats['leaks_detected'] += len(leaks)
                    self.last_action = f"Detected {len(leaks)} potential memory leaks"
                    
                    # Perform cleanup if leaks detected
                    cleanup_memory(aggressive=True)
                    self.stats['cleanups_performed'] += 1
                    
        except Exception as e:
            print(f"Leak detection error: {e}")
    
    def _create_status_report(self, enabled, config, current_usage):
        """Create management status report"""
        
        if not enabled:
            return "🔴 Smart Memory Management: DISABLED"
        
        # Status indicator
        if current_usage >= config['critical_threshold']:
            status_icon = "🔴"
            status_text = "CRITICAL"
        elif current_usage >= config['warning_threshold']:
            status_icon = "🟡"
            status_text = "WARNING"
        else:
            status_icon = "🟢"
            status_text = "NORMAL"
        
        uptime = time.time() - self.stats['uptime_start']
        uptime_str = f"{uptime/3600:.1f}h" if uptime > 3600 else f"{uptime/60:.1f}m"
        
        return (f"{status_icon} Smart Management: {status_text} | "
                f"Mode: {config['mode']} | "
                f"Memory: {current_usage:.1f}% | "
                f"Uptime: {uptime_str}")
    
    def _create_detailed_report(self, memory_info, config):
        """Create detailed management report"""
        
        report_lines = ["=== SMART MEMORY MANAGER REPORT ===\n"]
        
        # Management status
        report_lines.append(f"Management Status: {'ACTIVE' if self.management_active else 'INACTIVE'}")
        report_lines.append(f"Management Mode: {config['mode']}")
        report_lines.append(f"Last Action: {self.last_action}")
        
        # Thresholds
        report_lines.append(f"\nTHRESHOLDS:")
        report_lines.append(f"  Warning: {config['warning_threshold']:.1f}%")
        report_lines.append(f"  Critical: {config['critical_threshold']:.1f}%")
        report_lines.append(f"  Check Interval: {config['check_interval']:.1f}s")
        
        # Current memory state
        report_lines.append(f"\nCURRENT MEMORY STATE:")
        report_lines.append(f"  System Usage: {memory_info.get('system_percent', 0):.1f}%")
        report_lines.append(f"  System Total: {format_bytes(memory_info.get('system_total', 0))}")
        report_lines.append(f"  System Available: {format_bytes(memory_info.get('system_available', 0))}")
        report_lines.append(f"  Process RSS: {format_bytes(memory_info.get('process_rss', 0))}")
        
        # Features status
        report_lines.append(f"\nFEATURES:")
        report_lines.append(f"  Leak Detection: {'ENABLED' if config['leak_detection'] else 'DISABLED'}")
        report_lines.append(f"  VRAM Optimization: {'ENABLED' if config['vram_optimization'] else 'DISABLED'}")
        report_lines.append(f"  Aggressive Cleanup: {'ENABLED' if config['aggressive_cleanup'] else 'DISABLED'}")
        
        # Performance statistics
        uptime = time.time() - self.stats['uptime_start']
        report_lines.append(f"\nPERFORMANCE STATISTICS:")
        report_lines.append(f"  Uptime: {uptime/3600:.2f} hours")
        report_lines.append(f"  Total Cleanups: {self.stats['cleanups_performed']}")
        report_lines.append(f"  VRAM Optimizations: {self.stats['vram_optimizations']}")
        report_lines.append(f"  Memory Freed: {format_bytes(self.stats['memory_freed_total'])}")
        report_lines.append(f"  Leaks Detected: {self.stats['leaks_detected']}")
        
        # Efficiency metrics
        if uptime > 0:
            cleanup_rate = self.stats['cleanups_performed'] / (uptime / 3600)  # per hour
            report_lines.append(f"  Cleanup Rate: {cleanup_rate:.2f}/hour")
        
        return "\n".join(report_lines)
    
    def _create_performance_stats(self):
        """Create performance statistics summary"""
        
        uptime = time.time() - self.stats['uptime_start']
        
        if uptime < 60:
            uptime_str = f"{uptime:.0f}s"
        elif uptime < 3600:
            uptime_str = f"{uptime/60:.1f}m"
        else:
            uptime_str = f"{uptime/3600:.1f}h"
        
        return (f"Uptime: {uptime_str} | "
                f"Cleanups: {self.stats['cleanups_performed']} | "
                f"VRAM Opts: {self.stats['vram_optimizations']} | "
                f"Freed: {format_bytes(self.stats['memory_freed_total'])} | "
                f"Leaks: {self.stats['leaks_detected']}")
    
    def __del__(self):
        """Cleanup when node is destroyed"""
        try:
            if self.management_active:
                self._stop_smart_management()
        except:
            pass 